import torch
import numpy as np
import torch.nn as nn

# 随机创建一些训练数据
N, D_in, H, D_out = 64, 1000, 100, 10  # 输入，维度，中间层，输出
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)


class Net(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        super(Net, self).__init__()
        # define the model Architecture
        self.linear1 = torch.nn.Linear(D_in, H, bias=False)
        self.linear2 = torch.nn.Linear(H, D_out, bias=False)

    def forward(self, x):
        y_pred = self.linear2(self.linear1(x).clamp(min=0))
        return y_pred


# model
model = Net(D_in, H, D_out)

# loss function
loss_fn = torch.nn.MSELoss(reduction='sum')

# Optimizer
leaning_rate = 1e-4
optimizer = torch.optim.Adam(model.parameters(), lr=leaning_rate)

for it in range(500):
    # Forward pass
    y_pred = model(x)  # model.forward

    # compute loss
    loss = loss_fn(y_pred, y)  # MSE loss 均方误差
    print(it, loss.item())

    # Backward pass
    optimizer.zero_grad()
    loss.backward()

    # Optimizer - Update model parameter
    optimizer.step()
